//
// Created by f on 2021/7/8.
// 一个AVL树模板类的实现，代码简单易懂，没有做太多性能优化
//

#ifndef AVL_RBTREE_AVLTREE_H
#define AVL_RBTREE_AVLTREE_H

#include <algorithm>
#include <vector>
#include <stack>

using std::max;
using std::vector;
using std::stack;

template <typename T>
struct AVLnode{
    T value;
    int height;
    AVLnode * left;
    AVLnode * right;
};

template <typename T>
class AVLtree {
    typedef struct AVLnode<T> node_t;
    typedef struct AVLnode<T>* nodeptr_t;

private:
    nodeptr_t root_; //根节点
    int size_; //节点总数

    //计算树的高度
    int treeHeight(nodeptr_t root){
        if(root== nullptr) return 0;
        else return max(treeHeight(root->left), treeHeight(root->right))+1;
    }

    /*
     * 什么时候需要更新树的高度？
     * 1.在插入结点时，沿插入的路径更新结点的高度值
     * 2.在删除结点时（delete）,沿删除的路径更新结点的高度值
     */

    //计算平衡因子，也就是左子树和右子树的高度差
    //当平衡因子大于1时，触发树的再平衡
    int treeGetBalanceFactor(nodeptr_t root){
        if(root == nullptr) return 0;
        else{
            if(root->left == nullptr && root->right != nullptr) return -root->right->height;
            else if(root->left != nullptr && root->right == nullptr) return root->left->height;
            else if(root->left == nullptr && root->right == nullptr) return 0;
            else return root->left->height - root->right->height;
        }
    }

    /*
     * 树的再平衡由两个基本操作组成：左旋（逆时针）和右旋（顺时针）
     * 这种旋转在整个平衡化过程中可能进行一次或多次
     * 这两种操作都是从失去平衡的最小子树根结点开始的（即离插入结点最近且平衡因子超过1的祖结点）
     */

    //右旋
    nodeptr_t treeRotateRight(nodeptr_t root){
        nodeptr_t left = root->left;
        root->left = left->right; //左节点的右孩子，变成root的左孩子
        left->right = root; //root变成左节点的右孩子，左节点变成root
        //旋转后，左节点和root节点的孩子发生变化，需要更新高度
        left->height = max(treeHeight(left->left), treeHeight(left->right))+1;
        root->height = max(treeHeight(root->left), treeHeight(root->right))+1;

        return left; //旋转后左节点变成root了
    }

    //左旋，和左旋对称
    nodeptr_t treeRotateleft(nodeptr_t root){
        nodeptr_t right = root->right;
        root->right = right->left;
        right->left = root;

        right->height = max(treeHeight(right->left), treeHeight(right->right))+1;
        root->height = max(treeHeight(root->left), treeHeight(root->right))+1;

        return right;
    }

    /*
     * 需要进行再平衡的四种情况：
     * LL：在根节点左孩子的左子树添加节点，此时对根节点进行一次右旋
     * RR：在根节点右孩子的右子树添加节点，此时对根节点进行一次左旋
     * LR：在根节点左孩子的右子树添加节点，此时先左旋左孩子，再右旋根节点
     * RL：在根节点右孩子的左子树添加节点，此时先右旋右孩子，再左旋根节点
     * 如何判断属于那种情况？根据根节点及左右孩子的平衡因子判断
     */

    //再平衡
    nodeptr_t treeRebalance(nodeptr_t root){
        int factor = treeGetBalanceFactor(root);
        // LL
        if(factor>1 && treeGetBalanceFactor(root->left)>0)
            return treeRotateRight(root);
        // LR
        else if(factor>1 && treeGetBalanceFactor(root->left)<=0){
            root->left = treeRotateleft(root->left);
            return treeRotateRight(root);
        }
        // RR
        else if(factor<-1 && treeGetBalanceFactor(root->right)<=0)
            return treeRotateleft(root);
        // RL
        else if(factor<-1 && treeGetBalanceFactor(root->right)>0){
            root->right = treeRotateRight(root->right);
            return treeRotateleft(root);
        }
        // 本来就平衡，什么都不用干
        else return root;
    }

    // 删除树中值最小的节点，返回这个值
    T treeDeleteMin(nodeptr_t &root){
        //if(root == nullptr) return 0;

        //左孩子为空，则root就是最小值节点。把root清空，把右孩子替换为root
        //即使右孩子为空，也是运行这段代码
        if(root->left == nullptr){
            nodeptr_t right = root->right;
            T value = root->value;
            delete root;
            root = right;
            return value;
        }
            //左孩子不为空，递归左孩子
        else return treeDeleteMin(root->left);
    }

    //用于中序遍历的递归函数
    void tree2Vector(nodeptr_t root, vector<T> &ans){
        if(root==nullptr) return;
        tree2Vector(root->left, ans);
        ans.push_back(root->value);
        tree2Vector(root->right, ans);
    }

    //用于查找元素的递归函数
    nodeptr_t treeFind(nodeptr_t root, T value){
        if(root == nullptr) return nullptr;
        else{
            if(root->value == value) return root;
            else if(root->value < value) return treeFind(root->right, value);
            else return treeFind(root->left, value);
        }
    }

    //删除节点和子树
    void deleteSubtree(nodeptr_t root){
        if(root == nullptr) return;

        else{
            nodeptr_t left = root->left;
            nodeptr_t right = root->right;
            delete root;
            deleteSubtree(left);
            deleteSubtree(right);
        }
    }
    
    //用于插入元素的递归函数
    void treeInsert(nodeptr_t &root,T value){
        //树是空的，创建根节点
        if(root == nullptr){
            root = new node_t;
            root->value = value;
            root->right = nullptr;
            root->left = nullptr;
            root->height = 1;
            size_++;
            return;
        }
        if(root->value == value) return;
        else if(root->value < value){
            treeInsert(root->right, value);
            root->height = treeHeight(root); //递归结束，说明插入成功，回到父节点，更新父节点高度
        }
        else {
            treeInsert(root->left, value);
            root->height = treeHeight(root);
        }
        root = treeRebalance(root);
    }

    //用于删除元素的递归函数
    void treeDelete(nodeptr_t &root, T value){
        if(root == nullptr) return;

        if(root->value == value){
            //root有右子树，则用右子树中的最小值替换root的值，并删除这个值最小的节点
            if(root->right) root->value=treeDeleteMin(root->right);
                //root没有右子树，则用左孩子替换root，不管左孩子是不是空
            else{
                nodeptr_t left = root->left;
                delete root;
                root = left;
            }
            size_--;
        }
            //
        else if(root->value < value) treeDelete(root->right, value);
        else treeDelete(root->left, value);
        root = treeRebalance(root);
    }

    //用于层序遍历的递归函数
    void treeLevelOrder(nodeptr_t root, vector<vector<T>> &res, int depth){
        if(root==nullptr) return;
        if(depth>=res.size())
            res.push_back(vector<T> {});
        res[depth].push_back(root->value);
        treeLevelOrder(root->left, res, depth+1);
        treeLevelOrder(root->right, res, depth+1);
    }

    //用于前序遍历的递归函数
    void treePreOrder(nodeptr_t root, vector<T> &ans){
        if(root==nullptr) return;
        ans.push_back(root->value);
        treePreOrder(root->left, ans);
        treePreOrder(root->right, ans);
    }

    //用于后序遍历的递归函数
    void treePostOrder(nodeptr_t root, vector<T> &res) {
        if (root == nullptr) return;
        treePostOrder(root->left, res);
        treePostOrder(root->right, res);
        res.push_back(root->value);
    }

public:
    AVLtree(){
        root_ = nullptr;
        size_ = 0;
    }

    ~AVLtree(){
        deleteSubtree(root_);
    }

    // 删除值，然后对树进行再平衡
    void Delete(T value){
        treeDelete(root_, value);
    }

    // 插入值，然后对树进行再平衡
    // 每次插入都从根节点再平衡，这样并不好，还需要优化
    void Insert(T value){
        treeInsert(root_, value);
    }

    //TODO：实现树的迭代器。需要更改节点，保存指向父节点的指针，有点麻烦
    //基本思路是定义一个类中类，类成员包括指向节点的指针。重载加号和减号运算符

    //中序遍历，将AVL树转化为字典序列表
    vector<T> Vector() {
        vector<T> ans;
        tree2Vector(root_, ans);
        return ans;
    }

    //返回树的高度
    int Height(){
        return treeHeight(root_);
    }

    //返回树中节点的总数
    int Size(){
        return size_;
    }

    //在树中查找某个值，返回指向这个值的指针。如果找不到，返回空指针
    nodeptr_t Find(T value){
        nodeptr_t res = treeFind(root_, value);
        return res;
    }

    //输入值，返回索引。如果找不到就返回-1
    //用递归的方法不知道怎么中途跳出递归，导致每次都要做完整的中序遍历，这样实在糟糕
    //改用迭代+栈的方法
    int Value2Index(T value){
        stack<nodeptr_t> s;
        int index = 0;
        bool found = false;
        nodeptr_t p = root_;

        while(p!=nullptr || !s.empty()){
            if(p && p->value==value) {
                found = true;
                break;
            }
            while (p != nullptr) {
                s.push(p);
                p = p->left;
            }
            p = s.top();
            s.pop();
            index ++;
            p = p->right;
        }
        if(found) return index;
        else {
            cout<<"\033[31mERROR: VALUE NOT IN TREE \033[0m"<<endl;
            return -1;
        }
    }

    //输入索引，返回值。如果找不到就返回0且在终端报错
    //同样也是用迭代+栈的方法
    T At(int index){
        if(index<0 || index>=size_){
            cout<<"\033[31mERROR: INDEX OUT OF RANGE\033[0m"<<endl
            <<"\033[34mINDEX SHOULD BE IN RANGE \033[0m"<<"[0,"<<size_-1<<"]"<<endl;
            return 0;
        }

        stack<nodeptr_t> s;
        T value = 0;
        int i=0;
        nodeptr_t p = root_;

        while(p!=nullptr || !s.empty()){
            if(p!=nullptr && i==index) {
                value = p->value;
                break;
            }
            while (p != nullptr) {
                s.push(p);
                p = p->left;
            }
            p = s.top();
            s.pop();
            i++;
            p = p->right;
        }
        return index;
    }

    //层序遍历
    vector<vector<T>> LevelOrder() {
        vector<vector<T>> res;
        treeLevelOrder(root_, res, 0);
        return res;
    }

    //前序遍历
    vector<T> PreOrder() {
        vector<T> ans;
        treePreOrder(root_, ans);
        return ans;
    }

    //后序遍历
    vector<T> PostOrder() {
        vector<T> res;
        treePostOrder(root_, res);
        return res;
    }
};


#endif //AVL_RBTREE_AVLTREE_H
